package com.zyx.flinkdemo.dataset.stateprocessor;

import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.java.ExecutionEnvironment;
import org.apache.flink.api.java.operators.DataSource;
import org.apache.flink.api.scala.typeutils.Types;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.runtime.state.hashmap.HashMapStateBackend;
import org.apache.flink.state.api.BootstrapTransformation;
import org.apache.flink.state.api.OperatorTransformation;
import org.apache.flink.state.api.Savepoint;
import org.apache.flink.state.api.functions.StateBootstrapFunction;


/**
 * @author zyx
 * @since 2021/5/25 22:52
 * desc: StateProcessor写入实例
 */
public class StateProcessorWriteDemo {
    public static void main(String[] args) throws Exception {
        ExecutionEnvironment env = ExecutionEnvironment.getExecutionEnvironment();
        DataSource<Integer> input = env.fromElements(1, 2, 3, 4, 5, 6);
        // 将新的状态写入
        BootstrapTransformation<Integer> transformation = OperatorTransformation
                .bootstrapWith(input)
                .transform(new MySimpleBootstrapFunction());
        int maxParallelism = 128;
        Savepoint
                .create(new HashMapStateBackend(), maxParallelism)
                .withOperator("uid1", transformation)
                .write("hdfs://ns01/test/checkpoint_test/savepoint-1");

        env.execute();
    }

    private static class MySimpleBootstrapFunction extends StateBootstrapFunction<Integer> {

        private ListState<Integer> state;

        @Override
        public void processElement(Integer value, Context context) throws Exception {
            state.add(value);
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            state = context.getOperatorStateStore().getListState(new
                    ListStateDescriptor<>("state1", Types.INT()));
        }
    }
}
